# 将目标检测数据集分割为 验证集 和 训练集 和 测试集
import json
import os
import random

def tableDet():
    # 结构化文件标注的数据 作为test部分
    structured = 'I:/Images_OCR/structured/labelxx.json'
    allImages = 'I:/Images_OCR/cmrd_talbe/det/Voc_med/images'
    # 生成test文件
    test = 'I:/Images_OCR/cmrd_talbe/det/Voc_med/test.txt'
    valid = 'I:/Images_OCR/cmrd_talbe/det/Voc_med/valid.txt'
    train = 'I:/Images_OCR/cmrd_talbe/det/Voc_med/train.txt'
    with open(structured, 'r', encoding="utf-8") as f:
        structuredLines = json.load(f)
    structuredLists = []
    for ss in structuredLines:
        img_name = ss['img_name']
        structuredLists.append(img_name)
    all_images = os.listdir(allImages)

    testLists = []
    trainValidLists = []
    for img_name in all_images:
        s = 'images/' + img_name + ' ' + 'annotations/' + img_name.split('.')[0] + '.xml' + '\n'
        if img_name in structuredLists:
            testLists.append(s)
        else:
            trainValidLists.append(s)

    random.shuffle(trainValidLists)

    for i in range(len(testLists)):
        with open(test, 'a', encoding="utf-8") as file:
            file.write(testLists[i])

    for i in range(len(trainValidLists)):
        if i < len(trainValidLists) - len(testLists):
            with open(train, 'a', encoding="utf-8") as file:
                file.write(trainValidLists[i])
        else:
            with open(valid, 'a', encoding="utf-8") as file:
                file.write(trainValidLists[i])
    print('表格检测拆分成功')

def tableRec():
    # 结构化文件标注的数据 作为test部分
    structured = 'I:/Images_OCR/structured/labelxx.json'
    allImages = 'I:/Images_OCR/cmrd_talbe/rec/new/Voc_tableStructure/images'
    # 生成test文件
    test = 'I:/Images_OCR/cmrd_talbe/rec/new/Voc_tableStructure/test.txt'
    valid = 'I:/Images_OCR/cmrd_talbe/rec/new/Voc_tableStructure/valid.txt'
    train = 'I:/Images_OCR/cmrd_talbe/rec/new/Voc_tableStructure/train.txt'
    with open(structured, 'r', encoding="utf-8") as f:
        structuredLines = json.load(f)
    structuredLists = []
    for ss in structuredLines:
        img_name = ss['img_name']
        structuredLists.append(img_name)
    all_images = os.listdir(allImages)

    testLists = []
    trainValidLists = []
    for img_name in all_images:
        s = 'images/' + img_name + ' ' + 'annotations/' + img_name.split('.')[0] + '.xml' + '\n'
        if img_name in structuredLists:
            testLists.append(s)
        else:
            trainValidLists.append(s)

    random.shuffle(trainValidLists)

    for i in range(len(testLists)):
        with open(test, 'a', encoding="utf-8") as file:
            file.write(testLists[i])

    for i in range(len(trainValidLists)):
        if i < len(trainValidLists) - len(testLists):
            with open(train, 'a', encoding="utf-8") as file:
                file.write(trainValidLists[i])
        else:
            with open(valid, 'a', encoding="utf-8") as file:
                file.write(trainValidLists[i])
    print('表格结构识别拆分成功')

if __name__ == '__main__':
    # 表格检测
    tableDet()
    # 表格结构识别
    # tableRec()